#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
# @FileName  :mnist-train.py
# @Time      :2024/3/25 17:16
# @Author    :Toliong toliong@foxmail.com
"""
"""
代码解释

您提供的代码是一个使用PaddlePaddle框架的深度学习模型训练和预测流程的示例。这个示例使用了MNIST数据集来训练一个LeNet网络模型。以下是代码的详细解释：

1. **导入必要的库**：
   - `paddle`：PaddlePaddle深度学习框架。
   - `numpy`：用于处理数据的库。
   - `Normalize`：用于数据预处理，归一化数据。

2. **数据预处理**：
   - 使用`Normalize`函数对数据进行归一化处理，这里使用的均值和标准差都是127.5，数据格式为'CHW'，即通道-高度-宽度。

3. **数据集加载**：
   - 使用`paddle.vision.datasets.MNIST`下载并加载MNIST数据集，分为训练集和测试集，并应用之前定义的归一化处理。

4. **模型构建**：
   - 使用`paddle.vision.models.LeNet`构建LeNet网络模型，`num_classes=10`表示输出层有10个类别，对应MNIST数据集的10个数字。

5. **模型训练配置**：
   - 使用`paddle.Model`将LeNet模型封装成可训练的模型。
   - 配置优化器为Adam，损失函数为交叉熵损失，评价指标为准确率。

6. **模型训练**：
   - 使用`model.fit`方法对模型进行训练，指定训练数据集、迭代次数（epochs）、批量大小（batch_size）和是否打印详细信息（verbose）。

7. **模型评估**：
   - 使用`model.evaluate`方法评估模型在测试集上的性能，同样指定测试数据集和批量大小。

8. **模型保存与加载**：
   - 使用`model.save`保存训练好的模型到指定路径。
   - 使用`model.load`加载已保存的模型。

9. **模型预测**：
   - 从测试集中取出一张图片和对应的标签。
   - 调整图片的形状以匹配模型的输入要求。
   - 使用`model.predict_batch`进行预测，并使用`argmax`获取最可能的类别标签。

10. **结果展示**：
    - 打印出真实标签和预测标签。
    - 使用`matplotlib`库展示原始图片。
"""

import paddle
import numpy as np
from paddle.vision.transforms import Normalize
from matplotlib import pyplot as plt

transform = Normalize(mean=[127.5], std=[127.5], data_format='CHW')

# 下载数据集并初始化 DataSet
train_dataset = paddle.vision.datasets.MNIST(mode='train', transform=transform)
test_dataset = paddle.vision.datasets.MNIST(mode='test', transform=transform)

# 模型组网并初始化网络
lenet = paddle.vision.models.LeNet(num_classes=10)
model = paddle.Model(lenet)

# 模型训练的配置准备，准备损失函数，优化器和评价指标
model.prepare(paddle.optimizer.Adam(parameters=model.parameters()),
              paddle.nn.CrossEntropyLoss(),
              paddle.metric.Accuracy())

# 模型训练
model.fit(train_dataset, epochs=5, batch_size=64, verbose=1)
print('模型训练完成')
# 模型评估
model.evaluate(test_dataset, batch_size=64, verbose=1)
print('模型评估完成')

# 保存模型
model.save('./output/mnist')
# 加载模型
model.load('output/mnist')
print('保存完成')

# 从测试集中取出一张图片
img, label = test_dataset[0]
# 将图片shape从1*28*28变为1*1*28*28，增加一个batch维度，以匹配模型输入格式要求
img_batch = np.expand_dims(img.astype('float32'), axis=0)

# 执行推理并打印结果，此处predict_batch返回的是一个list，取出其中数据获得预测结果
out = model.predict_batch(img_batch)[0]
pred_label = out.argmax()
print('true label: {}, pred label: {}'.format(label[0], pred_label))

# 可视化图片
plt.imshow(img[0])
